import torch
import platform

class DeviceManager:
    """
    设备管理器类，负责GPU环境设置和设备管理
    """
    def __init__(self, config):
        self.config = config
        self.device = None
        self.dtype = None
        self.setup_gpu()
    
    def setup_gpu(self):
        """
        设置GPU环境，选择合适的设备和数据类型
        """
        try:
            # 检查CUDA可用性
            if torch.cuda.is_available():
                print("检测到NVIDIA GPU，将使用CUDA加速")
                device_id = self.config['gpu'].get('device_id', 0)
                
                # 检查指定的设备ID是否可用
                if device_id >= torch.cuda.device_count():
                    print(f"警告: 指定的设备ID {device_id} 不可用，将使用设备0")
                    device_id = 0
                    
                # 设置设备
                self.device = f"cuda:{device_id}"
                
                # 打印GPU信息
                gpu_name = torch.cuda.get_device_name(device_id)
                gpu_memory = torch.cuda.get_device_properties(device_id).total_memory / 1024**3  # 转换为GB
                print(f"使用GPU: {gpu_name} (设备{device_id})，总显存: {gpu_memory:.2f}GB")
                
                # 设置数据类型
                dtype_str = self.config['gpu'].get('dtype', 'bfloat16')
                if dtype_str == 'bfloat16':
                    if torch.cuda.is_bf16_supported():
                        self.dtype = torch.bfloat16
                        print("使用bfloat16数据类型")
                    else:
                        self.dtype = torch.float16
                        print("当前GPU不支持bfloat16，使用float16数据类型")
                elif dtype_str == 'float16':
                    self.dtype = torch.float16
                    print("使用float16数据类型")
                else:
                    self.dtype = torch.float32
                    print("使用float32数据类型")
                    
                # 设置最大内存
                max_memory = self.config['gpu'].get('max_memory', 'auto')
                if max_memory != 'auto':
                    try:
                        # 将GB转换为字节
                        if isinstance(max_memory, str) and 'GB' in max_memory:
                            max_memory_gb = float(max_memory.replace('GB', ''))
                            max_memory_bytes = int(max_memory_gb * 1024**3)
                        else:
                            max_memory_bytes = int(max_memory)
                        
                        # 设置CUDA内存限制
                        torch.cuda.set_per_process_memory_fraction(max_memory_bytes / torch.cuda.get_device_properties(device_id).total_memory, device_id)
                        print(f"设置GPU内存限制: {max_memory}")
                    except Exception as e:
                        print(f"设置内存限制时出错: {e}，将使用默认设置")
            else:
                # 检查MPS可用性（Mac）
                if platform.system().lower() == 'darwin' and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                    print("检测到Apple Silicon，将使用MPS加速")
                    self.device = "mps"
                    self.dtype = torch.float32  # MPS对float16和bfloat16的支持有限
                else:
                    print("未检测到GPU加速，将使用CPU")
                    self.device = "cpu"
                    self.dtype = torch.float32
            
            # 打印最终设备设置
            print(f"最终设备设置: device={self.device}, dtype={self.dtype}")
            
        except Exception as e:
            print(f"设置GPU环境时出错: {e}")
            # 出错时默认使用CPU
            self.device = "cpu"
            self.dtype = torch.float32
            print("已切换到CPU模式")
    
    def get_device(self):
        """
        获取当前设备
        """
        return self.device
    
    def get_dtype(self):
        """
        获取当前数据类型
        """
        return self.dtype
    
    def is_cuda_available(self):
        """
        检查是否使用CUDA
        """
        return self.device.startswith("cuda")
    
    def is_amp_enabled(self):
        """
        检查是否启用自动混合精度
        """
        return self.config['gpu'].get('enable_amp', True) and self.is_cuda_available()
    
    def get_gpu_info(self):
        """
        获取GPU信息
        """
        if self.is_cuda_available():
            device_id = int(self.device.split(":")[1])
            props = torch.cuda.get_device_properties(device_id)
            return {
                "name": props.name,
                "total_memory_gb": props.total_memory / 1024**3,
                "used_memory_gb": torch.cuda.memory_allocated(device_id) / 1024**3,
                "device_id": device_id
            }
        else:
            return None